﻿/*************************************************************************************
 * 工    具：  HK_UdpCommunicateSystem
 * 
 * 描    述：  该工具是为Unity设计的局域网UDP通信组件，支持单播、多播、广播三种通信模式，
 *             具备消息接收、ACK校验、多线程处理等特性，
 *             适用于实时通信场景（如多人VR、展厅多端交互）。
 * 
 * 版    本：  V1.5
 * 作    者：  京产肠饭
 * 
 * 创    建：  2024/08/27  V1.0
 * 
 * 更    新：  2024/10/16  V1.1
 *            1.代码优化：系统层代码优化，减少代码冗余，添加部分注释
 *            2.代码优化：指令层Command重新命名
 *            3.Bug修复：修复运行退出时，接收消息子线程的报错
 *
 *            2025/02/05  V1.2
 *            1.功能优化：显式配置广播（udpClient.EnableBroadcast = true），防止某些平台默认禁用广播
 *            2.代码优化：修复部分变量命名错误的问题
 *            3.功能优化：优化单例测试场景，通过UGUI实现了一个简单的测试场景，方便各平台测试
 *
 *            2025/05/03  V1.3
 *            1.代码优化：增加udpClient.Client.Shutdown(SocketShutdown.Both)，防止端口被占用
 *            2.代码优化：单例测试脚本中，UDP关闭调整到OnApplicationQuit方法中
 *
 *            2025/05/14  V1.4
 *            1.Bug修复：修复运行退出时，异常报错的BUG
 *
 *            2025/06/05  V1.5
 *            1.代码优化：点对点通信名称修改为 UnicastMsg 方法
 *            2.功能新增：新增 MulticastMsg，实现UDP多播消息发送
 *            
 * 链    接：  https://gitee.com/JingChanChangFan/hk_-unity-tools
*************************************************************************************/

using QFramework;
using System;
using System.Collections.Generic;
using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Threading;

namespace HKTools
{
    public interface IHK_UdpSystem : ISystem
    {
        /// <summary>
        /// 开启 UDP 通信
        /// </summary>
        void StartUDP(int localPort);

        /// <summary>
        /// 关闭 UDP 通信
        /// </summary>
        void CloseUDP();

        /// <summary>
        /// 单播消息
        /// </summary>
        void UnicastMsg(string msg, IPEndPoint remotePoint);

        /// <summary>
        /// 加入多播组
        /// </summary>
        void JoinMulticastGroup(IPAddress multicastAddress);

        /// <summary>
        /// 离开多播组
        /// </summary>
        void LeaveMulticastGroup(IPAddress multicastAddress);

        /// <summary>
        /// 多播消息
        /// </summary>
        /// <param name="msg">消息</param>
        /// <param name="multicastAddress">多播地址</param>
        /// <param name="multicastPort">多播端口号，若不填则默认发本机绑定端口</param>
        void MulticastMsg(string msg, IPAddress multicastAddress, int multicastPort = 0);

        /// <summary>
        /// 广播消息
        /// </summary>
        void BroadcastMsg(string msg, int broadcastPort);

        /// <summary>
        /// 初始化 关注的ACK消息列表，无需加ACK:标头
        /// </summary>
        void InitAckList(List<string> ackList);

        /// <summary>
        /// 广播消息（有校验）
        /// </summary>
        void BroadcastAckMsg(string msg, int broadcastPort, Action ensureCallback = null);
    }

    public class HK_UdpSystem : AbstractSystem, IHK_UdpSystem
    {
        // UDP通信是否正在运行
        volatile bool isRunning;

        UdpClient udpClient;

        // 用于在主线程中执行操作的上下文
        SynchronizationContext mainThreadContext;

        // 关注的 ACK 消息列表
        List<string> focusACKList;

        // 是否正在发送ACK消息
        volatile bool isSendingACK = false;
        // 是否收到ACK确认消息
        volatile bool isReceiveACK = false;
        // 存储最后发送的ACK消息
        volatile string lastACKMsg = "";

        int localPort = 0;

        // 记录已加入的多播组（避免重复加入）
        HashSet<IPAddress> joinedMulticastGroups = new HashSet<IPAddress>();

        protected override void OnInit()
        {
            isRunning = false;

            isReceiveACK = false;
            isSendingACK = false;

            mainThreadContext = SynchronizationContext.Current;

            focusACKList = new List<string>();

            HKDebug.Log("IHK_UdpSystem 初始化完成!", LogColor.Cyan, LogStyle.Bold);
        }

        /// <summary>
        /// 开启 UDP 通信
        /// </summary>
        public void StartUDP(int localPort)
        {
            // 如果当前是开启状态，无法再次开启
            if (isRunning)
            {
                HKDebug.LogWarning("UDP已启动，不能重复启动！");
                return;
            }

            // 创建UDP客户端并绑定本地端口
            udpClient = new UdpClient(localPort, AddressFamily.InterNetwork);

            this.localPort = localPort;

            // 启用UDP广播功能，允许发送广播消息到本地网络中的所有设备
            udpClient.EnableBroadcast = true;

            // 设置为运行状态
            isRunning = true;

            // 启动接收消息的线程
            ThreadPool.QueueUserWorkItem(ReceiveMsg);

            HKDebug.Log("开启UDP，端口号：" + localPort, LogColor.Blue);
        }

        // 接收消息的线程方法
        void ReceiveMsg(object obj)
        {
            IPEndPoint anyPoint;

            // 当UDP通信处于运行状态时
            while (isRunning)
            {
                try
                {
                    // 接收消息
                    anyPoint = new IPEndPoint(IPAddress.Any, 0);
                    byte[] bytes = udpClient.Receive(ref anyPoint);

                    string msg = Encoding.UTF8.GetString(bytes);

                    // 如果是 ACK 消息
                    if (msg.StartsWith("ACK:"))
                    {
                        // 去掉 ACK标头
                        string subMsg = msg.Substring(4);

                        bool isContains = false;

                        // 检查是否包含在关注列表中
                        foreach (string ack in focusACKList)
                        {
                            if (subMsg.Contains(ack) || string.Equals(subMsg, ack))
                            {
                                isContains = true;
                                break;
                            }
                        }

                        // 如果关注的ACK消息列表中包含该消息
                        if (isContains)
                        {
                            // 添加RACK标头
                            subMsg = "RACK:" + subMsg;

                            byte[] returnBytes = Encoding.UTF8.GetBytes(subMsg);

                            // 返回 ACK 确认消息
                            udpClient.Send(returnBytes, returnBytes.Length, anyPoint);

                            msg = msg.Substring(4);

                            HKDebug.Log("接收 ACK 消息：" + msg, LogColor.Blue);

                            mainThreadContext.Post(_ =>
                            {
                                this.SendEvent(new Event_UDP_ReceiveAckMsg() { _ackMsg = msg });
                            }, null);
                        }
                    }
                    else if (msg.StartsWith("RACK:")) // 如果是 RACK 消息
                    {
                        if (lastACKMsg == msg.Replace("RACK:", "ACK:"))
                        {
                            isReceiveACK = true;
                        }
                        HKDebug.Log("接收 RACK 消息：" + msg, LogColor.Blue);
                    }
                    else // 普通消息
                    {
                        HKDebug.Log("接收 消息：" + msg, LogColor.Blue);

                        mainThreadContext.Post(_ =>
                        {
                            this.SendEvent(new Event_UDP_ReceiveMsg() { _msg = msg });
                        }, null);
                    }
                }
                catch (SocketException se)
                {
                    if (isRunning)
                    {
                        // 过滤 ICMP 端口不可达/连接拒绝错误（常见错误码）
                        if (se.SocketErrorCode == SocketError.ConnectionRefused
                            || se.SocketErrorCode == SocketError.DestinationAddressRequired)
                        {
                            // 静默忽略（如发送到未监听的端口时，无需报错）
                            return;
                        }

                        // 处理Socket异常，比如接收超时或其他网络问题  
                        HKDebug.LogError("套接字异常：" + se.Message);
                    }
                }
                catch (Exception e)
                {
                    if (isRunning)
                        HKDebug.LogError("接收消息错误：" + e.Message);
                }
            }
        }

        /// <summary>
        /// 单播消息
        /// </summary>
        public void UnicastMsg(string msg, IPEndPoint remotePoint)
        {
            if (!isRunning || udpClient == null)
            {
                HKDebug.LogError("UDP未启动，无法发送单播消息！");
                return;
            }

            try
            {
                byte[] bytes = Encoding.UTF8.GetBytes(msg);
                udpClient.Send(bytes, bytes.Length, remotePoint);

                HKDebug.Log("发送 消息：" + msg, LogColor.Green);
            }
            catch (System.Exception e)
            {
                HKDebug.LogError("发送消息错误：" + e.Message);
            }
        }

        /// <summary>
        /// 加入多播组
        /// </summary>
        public void JoinMulticastGroup(IPAddress multicastAddress)
        {
            if (!isRunning || udpClient == null)
            {
                HKDebug.LogError("UDP未启动，无法加入多播组！");
                return;
            }

            if (joinedMulticastGroups.Contains(multicastAddress))
            {
                HKDebug.LogWarning($"已加入多播组 {multicastAddress}，无需重复加入");
                return;
            }

            try
            {
                // 加入多播组（默认使用本地IP，可根据需求扩展）
                udpClient.JoinMulticastGroup(multicastAddress);
                joinedMulticastGroups.Add(multicastAddress);
                HKDebug.Log($"加入多播组 {multicastAddress} 成功", LogColor.Green);

                // 可选：启用/禁用多播消息本地回环（默认启用）
                udpClient.MulticastLoopback = true; // 自己发送的多播消息会被本地接收
            }
            catch (Exception e)
            {
                HKDebug.LogError($"加入多播组 {multicastAddress} 失败：{e.Message}");
            }
        }

        /// <summary>
        /// 离开多播组
        /// </summary>
        public void LeaveMulticastGroup(IPAddress multicastAddress)
        {
            if (!isRunning || udpClient == null)
            {
                HKDebug.LogError("UDP未启动，无法离开多播组！");
                return;
            }

            if (!joinedMulticastGroups.Contains(multicastAddress))
            {
                HKDebug.LogWarning($"未加入多播组 {multicastAddress}，无需离开");
                return;
            }

            try
            {
                udpClient.DropMulticastGroup(multicastAddress);
                joinedMulticastGroups.Remove(multicastAddress);
                HKDebug.Log($"离开多播组 {multicastAddress}: 成功", LogColor.Green);
            }
            catch (Exception e)
            {
                HKDebug.LogError($"离开多播组 {multicastAddress}: 失败：{e.Message}");
            }
        }

        /// <summary>
        /// 多播消息
        /// </summary>
        public void MulticastMsg(string msg, IPAddress multicastAddress, int multicastPort = 0)
        {
            if (!isRunning || udpClient == null)
            {
                HKDebug.LogError("UDP未启动，无法发送多播消息！");
                return;
            }

            if (!joinedMulticastGroups.Contains(multicastAddress))
            {
                if (multicastPort != 0)
                    HKDebug.LogError($"未加入多播组 {multicastAddress}:{multicastPort}，无法发送消息！");
                else
                    HKDebug.LogError($"未加入多播组 {multicastAddress}:{localPort}，无法发送消息！");
                return;
            }

            try
            {
                IPEndPoint multicastPoint = null;

                if (multicastPort != 0)
                    multicastPoint = new IPEndPoint(multicastAddress, multicastPort);
                else
                    multicastPoint = new IPEndPoint(multicastAddress, localPort);

                byte[] bytes = Encoding.UTF8.GetBytes(msg);

                if (multicastPort != 0)
                    HKDebug.Log($"多播 消息 ({multicastAddress}:{multicastPort})： {msg}", LogColor.Green);
                else
                    HKDebug.Log($"多播 消息 ({multicastAddress}:{localPort})： {msg}", LogColor.Green);

                udpClient.Send(bytes, bytes.Length, multicastPoint);
            }
            catch (Exception e)
            {
                HKDebug.LogError($"发送多播消息失败：{e.Message}");
            }
        }

        /// <summary>
        /// 广播消息
        /// </summary>
        public void BroadcastMsg(string msg, int broadcastPort)
        {
            if (!isRunning || udpClient == null)
            {
                HKDebug.LogError("UDP未启动，无法发送广播消息！");
                return;
            }

            try
            {
                HKDebug.Log("广播 消息：" + msg, LogColor.Green);

                IPEndPoint broadcastPoint = new IPEndPoint(IPAddress.Broadcast, broadcastPort);

                byte[] bytes = Encoding.UTF8.GetBytes(msg);
                udpClient.Send(bytes, bytes.Length, broadcastPoint);
            }
            catch (System.Exception e)
            {
                HKDebug.LogError("广播消息错误：" + e.Message);
            }
        }

        /// <summary>
        /// 初始化 关注的ACK消息列表，无需加ACK:标头
        /// </summary>
        public void InitAckList(List<string> ackList)
        {
            ackList.ForEach(ackMsg => focusACKList.Add(ackMsg));
        }

        /// <summary>
        /// 广播消息（有校验）
        /// </summary>
        public void BroadcastAckMsg(string msg, int broadcastPort, Action ensureCallback = null)
        {
            if (!isRunning)
                return;

            if (isSendingACK)
            {
                HKDebug.LogWarning("上一个ACK指令未确认，无法再次发送！");
                return;
            }

            ThreadPool.QueueUserWorkItem(_ =>
            {
                isReceiveACK = false; // 默认未收到 返回的 RACK
                isSendingACK = true; // 正在发送中

                int tryNum = 0;

                // 需要确认的消息，开头需加上ACK标头
                msg = "ACK:" + msg;

                lastACKMsg = msg;

                HKDebug.Log("广播 ACK消息：" + msg, LogColor.Green);

                // 最大尝试次数
                int maxTryNum = 5;

                while (tryNum < maxTryNum && !isReceiveACK)
                {
                    // 发送 ACK 广播消息
                    IPEndPoint broadcastPoint = new IPEndPoint(IPAddress.Broadcast, broadcastPort);

                    byte[] bytes = Encoding.UTF8.GetBytes(msg);
                    udpClient.Send(bytes, bytes.Length, broadcastPoint);

                    Thread.Sleep(250);

                    if (isReceiveACK)
                    {
                        HKDebug.Log("收到 ACK确认消息！", LogColor.Green);
                        isSendingACK = false;

                        mainThreadContext.Post(_ =>
                        {
                            ensureCallback?.Invoke();
                        }, null);

                        break;
                    }

                    HKDebug.LogWarning("未收到 ACK确认消息，重复次数：" + tryNum);
                    tryNum++;
                }

                isSendingACK = false;
            });
        }

        /// <summary>
        /// 关闭 UDP 通信
        /// </summary>
        public void CloseUDP()
        {
            if (!isRunning || udpClient == null)
            {
                HKDebug.LogWarning("UDP未启动！");
                return;
            }

            isRunning = false;

            // 离开所有已加入的多播组
            foreach (var addr in joinedMulticastGroups)
            {
                try { udpClient.DropMulticastGroup(addr); }
                catch { /* 忽略异常 */ }
            }
            joinedMulticastGroups.Clear();

            try
            {
                if (udpClient != null)
                {
                    // 关闭发送和接收，但不释放资源（避免重复关闭）
                    // 注意：若套接字已关闭，此操作可能抛出异常，需捕获
                    try
                    {
                        if (udpClient.Client != null && udpClient.Client.Connected)
                        {
                            udpClient.Client.Shutdown(SocketShutdown.Both);
                        }
                    }
                    catch (Exception innerEx)
                    {
                        HKDebug.LogWarning("UDP套接字关闭时状态异常：" + innerEx.Message);
                    }

                    // 释放资源（包含Close()的逻辑）
                    udpClient.Dispose();
                }
            }
            catch (Exception e)
            {
                HKDebug.LogError("关闭UDP时发生错误：" + e.Message);
            }
            finally
            {
                udpClient = null;
            }

            HKDebug.Log("关闭UDP", LogColor.Blue);
        }
    }
}